import pandas as pd
from pylab import *

# 设置中文字体和关闭警告
mpl.rcParams['font.family'] = ['Times New Roman', 'SimSun']
mpl.rcParams['axes.unicode_minus'] = False     # 解决负号显示问题
mpl.rcParams['font.size'] = 14  # 增大基础字体大小

# 读取数据
df = pd.read_csv('rescue_schedule_by_team.csv', encoding='cp1252')
df.columns = ['时间', '1', '2', '3']

# 创建图形
figure(figsize=(12, 5))  # 减小高度以使方格变窄
xlim(24, 96)
xticks(arange(24, 97, 12))
ylim(0.5, 3.5)  # 调整y轴范围，让间距更紧凑
yticks([1, 2, 3], ['1', '2', '3'])  # 明确设置刻度位置和标签

def merge_continuous_sections(team_data):
    merged = []
    if team_data.empty:
        return merged
    times = team_data['时间'].values
    sections = team_data[team_data.columns[1]].values
    current_start = times[0]
    current_end = times[0]
    current_section = sections[0]

    for i in range(1, len(times)):
        if times[i] == current_end + 1 and sections[i] == current_section:
            current_end = times[i]
        else:
            merged.append((current_start, current_end, current_section))
            current_start = current_end = times[i]
            current_section = sections[i]
    merged.append((current_start, current_end, current_section))
    return merged

# 使用白色替代之前的颜色方案
for idx, team in enumerate(['1', '2', '3']):
    team_data = df[['时间', team]].dropna().sort_values('时间')
    sections = merge_continuous_sections(team_data)

    for start, end, section in sections:
        duration = end - start + 1
        barh(float(team), duration, left=start, color='white', edgecolor='black', height=0.2)  # 减小高度值
        # 在路段编号后添加计算出的工作时间
        work_time = duration  # 计算工作时间
        text(start + duration/2, float(team),  # 确保team被转换为数值
             f'{section},{work_time}',
             ha='center', va='center', fontsize=14)

xlabel('时间/h', fontsize=16)
ylabel('修复施工队编号', fontsize=16)
tick_params(axis='both', which='major', labelsize=14)
grid(True, linestyle='--', alpha=0.7)
gca().spines['top'].set_visible(False)
gca().spines['right'].set_visible(False)
tight_layout()
savefig('rescue_schedule_gantt.png', dpi=300, bbox_inches='tight')
show()